package com.yihaodian.search.broker;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

import org.apache.log4j.Logger;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.inject.Inject;
import com.yihaodian.api.client.BrokerService;
import com.yihaodian.api.common.Request;
import com.yihaodian.api.common.Response;
import com.yihaodian.api.index.IndexRequest;
import com.yihaodian.api.search.SearchException;
import com.yihaodian.api.search.SearchRequest;
import com.yihaodian.api.search.SearchResponse;
import com.yihaodian.api.search.SugRequest;
import com.yihaodian.api.search.SugResponse;
import com.yihaodian.api.search.SugResponse.SugResult;
import com.yihaodian.cluster.ClusterService;
import com.yihaodian.cluster.ClusterState;
import com.yihaodian.cluster.Node;
import com.yihaodian.index.meta.routing.RoutingTable;
import com.yihaodian.index.meta.shard.Shard;
import com.yihaodian.search.broker.shard.ShardIndexClient;
import com.yihaodian.search.broker.shard.ShardIndexRequest;
import com.yihaodian.search.broker.shard.ShardSearchClient;
import com.yihaodian.search.broker.shard.ShardSearchRequest;
import com.yihaodian.search.broker.shard.ShardSearchResponse;

public class SearchBroker implements BrokerService {
	private static Logger log = Logger.getLogger(SearchBroker.class);

	public static long shard_timeout = 3000; // 3 seconds

	public static TimeUnit shard_timeout_unit = TimeUnit.MILLISECONDS;

	protected ClusterService clusterService;
	protected ShardSearchClient searchClient;
	protected ShardIndexClient indexClient;
	protected OperationRouting operationRouting;

	@Inject
	public SearchBroker(ClusterService clusterService,
			ShardSearchClient client, ShardIndexClient indexClient) {
		this.clusterService = clusterService;
		this.operationRouting = new OperationRouting();
		this.searchClient = client;
		this.indexClient = indexClient;
	}

	@Override
	public Response execute(Request request) throws SearchException {

		Response response = null;
		if (request instanceof SearchRequest) {
			response = search((SearchRequest) request);
		} else if (request instanceof SugRequest) {
			response = suggest((SugRequest) request);
		} else if (request instanceof IndexRequest) {
			index((IndexRequest) request);
		}

		return response;
	}

	/**
	 * scatter and gather (merge).
	 */
	public SearchResponse search(final SearchRequest request) {
		ShardRespCollector collector = scatter(request);

		Map<String, Response> respMap = collector.shardResponses;
		List<ShardSearchResponse> ssr = Lists.newArrayList();
		for (Response s : respMap.values())
			ssr.add((ShardSearchResponse) s);

		SearchResponse response = ResultMerger.merge(request, ssr);

		response.setNode(clusterService.getLocalNode().getId());

		return response;
	}

	public SugResponse suggest(SugRequest request) {
		ShardRespCollector collector = scatter(request);

		Map<String, SugResult> resultMap = Maps.newHashMap();
		Map<String, Response> respMap = collector.shardResponses;
		for (Response s : respMap.values()) {
			SugResponse sr = (SugResponse) s;

			for (SugResult sug : sr.getResult()) {
				SugResult ret = resultMap.get(sug.key);
				if (ret == null) {
					resultMap.put(sug.key, sug);
				} else {
					ret.value = ret.value + sug.value;
				}
			}
		}

		List<SugResult> results = new ArrayList<SugResult>(resultMap.values());
		Collections.sort(results, Collections.reverseOrder());
		if (request.getNumber() < results.size())
			results = results.subList(0, request.getNumber());

		SugResponse response = new SugResponse();
		response.setResult(results);

		response.setNode(clusterService.getLocalNode().getId());

		return response;
	}

	ShardRespCollector scatter(Request request) {
		if (clusterService.state() == null)
			throw new SearchException("ClusterState is not init yet");

		List<Shard> shards = operationRouting.shardRequest(
				clusterService.state(), request.getIndexName());

		ShardRespCollector collector = new ShardRespCollector(shards.size());
		for (Shard shard : shards) {
			submit(shard.getShardId().getShardId(), request, collector);
		}

		try {
			// wait 3 seconds for single shard
			collector.latch.await(shard_timeout, shard_timeout_unit);
		} catch (InterruptedException e) {
			log.error(e.getMessage(), e);
		}

		return collector;
	}

	/**
	 * usually Send request to target Node. We make use of hedwig to do load
	 * balance, so ShardSearchRequest does not contain <Node>.
	 * 
	 * @param request
	 * @param respCollector
	 */
	void submit(final String shardId, final Request request,
			final ShardRespCollector respCollector) {
		ExecutorService es = ThreadPools.getThreadPool(ThreadPools.SEARCH);

		es.submit(new Runnable() {
			@Override
			public void run() {
				try {
					Response response = null;

					if (request instanceof SearchRequest) {
						ShardSearchRequest sr = new ShardSearchRequest(shardId,
								(SearchRequest) request);
						response = searchClient.executeSearch(sr);
					} else if (request instanceof SugRequest) {
						response = searchClient.executeSug(
								request.getIndexName(), shardId,
								(SugRequest) request);
					}

					respCollector.collect(shardId, response);
				} catch (Exception e) {
					log.error(e.getMessage(), e);

					respCollector.latch.countDown();
				}
			}
		});
	}

	static class ShardRespCollector {
		CountDownLatch latch = null;
		Map<String, Response> shardResponses = Maps.newHashMap();

		public ShardRespCollector(int shardNum) {
			latch = new CountDownLatch(shardNum);
		}

		public void collect(String shardId, Response resp) {
			if (resp != null)
				shardResponses.put(shardId, resp);
			latch.countDown();
		}
	}

	public void index(IndexRequest request) throws SearchException {
		try {
			List<ShardIndexRequest> shardRequests = operationRouting
					.shardRequest(clusterService.state(), request);

			ClusterState cs = clusterService.state();
			String indexId = request.getIndexName();
			RoutingTable table = cs.routingTables().routingTable(indexId);
			for (ShardIndexRequest shardReq : shardRequests) {
				String shardId = shardReq.getShardId();
				Shard shard = new Shard(indexId, shardId);
				Collection<String> nodes = table.getNodes(shard.getShardId());

				for (String nodeId : nodes) {
					Node node = new Node(nodeId);
					submit(node, shardReq);
				}
			}

		} catch (Exception e) {
			throw new SearchException(e);
		}
	}

	void submit(final Node node, final ShardIndexRequest request) {
		// log.debug("submit index request, request:  " + request + ", node: " +
		// node.getId());

		ExecutorService es = ThreadPools.getThreadPool(ThreadPools.INDEX);

		es.submit(new Runnable() {

			@Override
			public void run() {
				try {
					indexClient.index(node.getIp(), node.getPort(), request);
				} catch (Exception e) {
					log.error("index error " + node, e);
				}
			}
		});
	}
}
